"""
@Author  : 吕申凯
@Time    : 2022/9/18 19:50
@File    : Group_normalization.py
@Function: 
"""

import torch
import torch.nn as nn
from tools.common_tools import set_seed

if __name__ == '__main__':
    set_seed(1)  # 设置随机种子

    batch_size = 2
    num_features = 4  # 特征数
    num_groups = 2  # 分组数必须被特征数整除

    features_shape = (2, 2)

    feature_map = torch.ones(features_shape)  # 2D
    feature_maps = torch.stack([feature_map * (i + 1) for i in range(num_features)], dim=0)  # 3D
    feature_maps_bs = torch.stack([feature_maps * (i + 1) for i in range(batch_size)], dim=0)  # 4D

    gn = nn.GroupNorm(num_groups, num_features)
    outputs = gn(feature_maps_bs)

    print("Group Normalization")
    print(gn.weight.shape)
    print(outputs[0])
